#!/bin/bash
# applicable to Ascend

echo "=============================================================================================================="
echo "Please run the script as: "
echo "bash run.sh RANK_SIZE"
echo "For example: bash run.sh /path/dataset 8"
echo "It is better to use the absolute path."
echo "=============================================================================================================="
set -e

source /home/liupan/.bashrc

RANK_SIZE=$1
ROOT_FOLDER="/home/liupan/WingsofPanda/WP-GeneralPurposeMedicalFoundationModels/wpgpmfm"
#
#EXEC_PATH=$(pwd)
#PYTHONPATH="/home/liuyu/code/wpgeneralpurposemedicalfoundationmodels"

#test_dist_8pcs()
#{
#    export RANK_TABLE_FILE=${ROOT_FOLDER}/unet3d/rank_table_8devices.json
#    export RANK_SIZE=8
#}
#
#test_dist_4pcs()
#{
#    export RANK_TABLE_FILE=${ROOT_FOLDER}/unet3d/rank_table_4devices.json
#    export RANK_SIZE=4
#}

test_dist_2pcs()
{
    export RANK_TABLE_FILE=${ROOT_FOLDER}/unet3d/rank_table_2devices.json
    export RANK_SIZE=2
}

test_dist_${RANK_SIZE}pcs

for((i=0;i<${RANK_SIZE};i++))
do
    rm -rf device${i+2}
    mkdir device${i+2}

    export DEVICE_ID=${i+2}
    export RANK_ID=${i}
    echo "start training for device ${i+2}"
    env > ./device${i+2}/env.log
    # silence this process
    PYTHONPATH=${PYTHONPATH} python3 ${ROOT_FOLDER}/unet3d/main_parallal.py > ./device${i+2}/train.log 2>&1 &
done
rm -rf device0
mkdir device0

export DEVICE_ID=2
export RANK_ID=0
echo "start training for device 0"
env > ./device0/env.log
PYTHONPATH=${PYTHONPATH} python3 ${ROOT_FOLDER}/unet3d/main_parallal.py 2>&1 | tee ./device0/train.log
if [ $? -eq 0 ];then
    echo "training success"
else
    echo "training failed"
    exit 2
fi
